# -*- coding: utf-8 -*-
"""
Created on Tue Jun 26 11:29:13 2012

@author: nzhao
"""
from traits.api import HasTraits, Button, Float, Range, Array, Instance, Int, List
from traitsui.api import View, Group, HGroup, Item, RangeEditor, ArrayEditor, InstanceEditor
from enable.api import Component, ComponentEditor
from chaco.api import VPlotContainer, HPlotContainer

import numpy as np

from math_qm.CoordianteSystem import CoordinateSystem
from phy.spin.MultiSpin import MultiSpin
from phy.control.CPMG import CPMG
from gui.cpmg_spin_half import cpmg_spin_half, control_group
#from misc.multilines import OverlappingPlotContainer

C13bath_group= Group(HGroup(
                            Item('do_calc', show_label=False),
                            Item('cut_off_dist', show_label=False),
                            Item('spin_include', show_label=False),
                        ),
                     control_group,
                     Item('plot123', editor=ComponentEditor(), show_label=False),
                     label = 'C-13 independent bath'
                    )

class C13Bath(cpmg_spin_half, HasTraits):
    do_calc      = Button
    coh_list     = Array()
    cut_off_dist = Float(100.0)
    spin_include = Int(1)
    multi_spin   = Instance(MultiSpin)
    
    cpmg_calculator = Instance(CPMG)
    
    axes_b = Instance( CoordinateSystem )
    axes_list   = List( CoordinateSystem )

    def __init__(self,  multi_spin = None, axes = None):
        super(C13Bath, self).__init__()
#        if multi_spin is not None:
#            self.multi_spin = multi_spin
#            self.multi_spin.on_trait_change(self.update_maxN, 'spin_num')
#        else:
#            self.multi_spin = MultiSpin()
#            
#        if axes is not None:
#            self.axes_list = axes
#            self.axes_b = axes[0]
#        else:
#            self.axes_b = CoordinateSystem()
#                    
#        self.is_auto_update = False
#        self.plot1=OverlappingPlotContainer(padding=30, fill_padding=True, bgcolor="white", use_backbuffer=True)
#        self.plot2=OverlappingPlotContainer(padding=30, fill_padding=True, bgcolor="white", use_backbuffer=True)
#        self.plot3=OverlappingPlotContainer(padding=30, fill_padding=True, bgcolor="white", use_backbuffer=True)
#        self.cpmg_calculator = CPMG(spin_num = 1)
   
        
    def isCalc(self, spin):
        res = (spin.dist < self.cut_off_dist) and (spin.spin_mark == 'None')
        return res
    
    def _cut_off_dist_changed(self):
        ncount = 0
        for spin in self.multi_spin.spin_list:
            if self.isCalc(spin):
                ncount = ncount + 1
        self.spin_include = ncount
        
    def _do_calc_fired(self):
        self.is_auto_update = False

        xlist = eval(self.time_list_str)
        self.coh_list = np.ones(len(xlist))
        
        for spin in self.multi_spin.spin_list:
            if spin.spin_type is 'Bath' and self.isCalc(spin):
                self.Ax, self.Ay, self.Az = self.multi_spin.hyperfine(spin, self.axes_b)
                self.update_fields()
                
                state_vec = np.array([spin.state_X, spin.state_Y, spin.state_Z])
                
                ylist = self.calculate_data(xlist)
                complex_coh = np.array([vec[0] + 1.j * np.dot(state_vec, vec[1:4]) for vec in ylist[:, 0]])
                
                self.coh_list = self.coh_list * complex_coh
        self.plot_data()
        
    def plot_data(self):
        xlist = eval(self.time_list_str)
        
        self.plot1.clear_data()
        self.plot2.clear_data()
        self.plot3.clear_data()

        self.plot1.add_data([xlist, self.coh_list.real])
        self.plot1.add_data([xlist, self.coh_list.imag])
        self.plot2.add_data([xlist, self.coh_list.real])
        self.plot2.add_data([xlist, self.coh_list.imag])
        self.plot3.add_data([xlist, self.coh_list.real])
        self.plot3.add_data([xlist, self.coh_list.imag])

        self.plot1.plot_data(['black', 'red'])
        self.plot2.plot_data(['black', 'red'])
        self.plot3.plot_data(['black', 'red'])
        
        self.plot23  = HPlotContainer(self.plot2, self.plot3, stack_order="left_to_right")
        self.plot123 = VPlotContainer(self.plot1, self.plot23, stack_order="top_to_bottom")
        self.plot2.components[0].index =self.plot1.components[0].index
        self.plot3.components[0].index =self.plot1.components[0].index
        
    traits_view = View(
        C13bath_group,
        title     = 'C13 bath',
        width     = .8,
        height    = .9,
        resizable = True,
        buttons   = [ 'OK' ],
        kind      = 'live'
    )

if __name__ == '__main__':
    bath = C13Bath()
    bath.configure_traits()